import torch
import torch.nn as nn
from modules.vit_seg_modeling import VisionTransformer as transUnet
from modules.vit_seg_modeling import CONFIGS as CONFIGS_ViT_seg

# 封装为攻击模块
class TransUNet_Attack(nn.Module):
    def __init__(self, residual=True, attack=True):
        """
        residual: 如果为 True，则返回 image + noise（残差扰动）；
                  否则直接返回生成后的攻击图像
        attack: 控制是否启用攻击模块
        """
        super(TransUNet_Attack, self).__init__()
        self.attack = attack
        self.residual = residual
        vit_name='R50-ViT-B_16'
        config_vit = CONFIGS_ViT_seg[vit_name]
        config_vit.n_classes = 3
        config_vit.n_skip = 3
        if vit_name.find('R50') != -1:
            config_vit.patches.grid = (int(128 / 16), int(128 / 16))
        self.unet = transUnet(config_vit, img_size=128, num_classes=config_vit.n_classes).cuda()
 
    def forward(self, image):
        if not self.attack:
            return image

        noise_or_image = self.unet(image)

        if self.residual:
            attacked = torch.clamp(image + noise_or_image, 0, 1)
        else:
            attacked = noise_or_image

        return attacked


